#pragma once
#include <vector>
#include <map>
#include <string>
#include <torch/torch.h>
#include "SphericalHarmonic.h"

struct DCBResult {
    std::string id;         // 卫星或接收机ID
    double dcb;             // DCB值(ns)
    double dcbSigma;        // DCB标准差
};

class DCBEstimator {
public:
    DCBEstimator();
    ~DCBEstimator();

    // 设置观测数据
    void setObservations(const std::vector<VTECObservation>& vtecObs);
    
    // 设置球谐函数模型
    void setSphericalHarmonicModel(const SphericalHarmonic& shModel);
    
    // 估计DCB（使用先前拟合的球谐模型）
    void estimateDCB();
    
    // 联合估计VTEC球谐系数和DCB
    void estimateVTECAndDCB(int maxDegree);
    
    // 获取卫星DCB结果
    std::vector<DCBResult> getSatelliteDCBs() const;
    
    // 获取接收机DCB结果
    std::vector<DCBResult> getReceiverDCBs() const;
    
    // 获取卫星DCB结果(字符串到double的映射)
    std::map<std::string, double> getSatelliteDCBMap() const;
    
    // 获取接收机DCB结果(字符串到double的映射)
    std::map<std::string, double> getReceiverDCBMap() const;
    
    // 获取估计的球谐系数
    std::vector<double> getSphericalHarmonicCoefficients() const;

    // Add this method
    void setZeroMeanConstraint(bool enable) {
        useZeroMeanConstraint_ = enable;
    }

private:
    std::vector<VTECObservation> vtecObservations_;
    SphericalHarmonic shModel_;
    std::vector<DCBResult> satelliteDCBs_;
    std::vector<DCBResult> receiverDCBs_;
    std::set<std::string> satellites_;
    std::set<std::string> receivers_;
    std::vector<double> shCoefficients_; // 联合估计的球谐系数
    
    // 构建DCB估计的设计矩阵
    torch::Tensor buildDCBDesignMatrix() const;
    
    // 构建联合估计的设计矩阵
    torch::Tensor buildJointDesignMatrix(int maxDegree) const;
    
    // 应用DCB零均值约束
    void applyZeroMeanConstraint(torch::Tensor& designMatrix, torch::Tensor& observationVector);

    // 从解向量中提取DCB值和球谐系数
    void extractDCBAndSHValues(const torch::Tensor& solution, int numSHCoeffs);
    
    // 从解向量中提取DCB值
    void extractDCBValues(const torch::Tensor& solution);

    bool useZeroMeanConstraint_ = true;
};